#include "StdAfx.h"
#include "Socket.h"
#include "SocketException.h"

#define WSA_CALL(fun) if(fun == SOCKET_ERROR) \
	throw SocketException(WSAGetLastError(), #fun);

Socket::Socket(Socket::SocketType type, Socket::SocketFamily fam, Socket::SocketProtocol proto) : 
mWsaCreated(NULL),
mFamily(fam),
mType(type),
mProtocol(proto) {
	mRemotePoint.isIPv6 = false;
	mRemotePoint.IPv4 = 0xFFFFFFFF;

	int32 retVal = WSAStartup(MAKEWORD(2, 0), &mWsdData);
	if(retVal != 0)
		throw SocketException(retVal);
	
	mWsaCreated = true;
	mSocket = socket((int)fam, (int)type, (int)proto);
	if(mSocket == INVALID_SOCKET)
		throw SocketException(WSAGetLastError());
}

Socket::Socket(SOCKET s, SocketType type, SocketFamily family, SocketProtocol proto) : 
mSocket(s),
mWsaCreated(false),
mFamily(family),
mType(type),
mProtocol(proto) {
	
}

Socket::~Socket() {
	try { disconnect(); } catch(...) { }

	if(mWsaCreated == true)
		WSACleanup();
}

Socket::HostEntry Socket::resolveHost(const std::string& hostOrIp) {
	ADDRINFO hint = { 0 };
	hint.ai_family = (int)mFamily;
	hint.ai_protocol = (int)mProtocol;
	hint.ai_socktype = (int)mType;

	PADDRINFOA resultArray = NULL;
	int32 error = getaddrinfo(hostOrIp.c_str(), NULL, &hint, &resultArray);
	if(error != 0)
		throw SocketException(error);

	PADDRINFOA curInfo = resultArray;
	while(curInfo != NULL) {
		if(curInfo->ai_family == (int)mFamily) {
			HostEntry ret;
			ret.isIPv6 = mFamily == IPv6;
			if(ret.isIPv6) {
				sockaddr_in6* sin6 = (sockaddr_in6*)curInfo->ai_addr;
				memcpy(&ret.IPv6, sin6->sin6_addr.u.Byte, 16);
			} else {
				sockaddr_in* sin = (sockaddr_in*)curInfo->ai_addr;
				ret.IPv4 = sin->sin_addr.s_addr;
			}

			freeaddrinfo(resultArray);
			return ret;
		}
		curInfo = curInfo->ai_next;
	}

	if(resultArray != NULL)
		freeaddrinfo(resultArray);

	throw SocketException(WSAHOST_NOT_FOUND);
}

void Socket::connect(const std::string& hostOrIp, uint16 port) {
	auto addr = resolveHost(hostOrIp);
	union
	{
		sockaddr_in addr4;
		sockaddr_in6 addr6;
	};
	sockaddr* socketAddr = NULL;
	uint32 addrSize = 0;
	switch(mFamily) {
	case IPv4: 
		{
			addrSize = sizeof(sockaddr_in);
			addr4.sin_addr.s_addr = addr.IPv4;
			addr4.sin_port = htons(port);
			socketAddr = (sockaddr*)&addr4;
		}
		break;
	case IPv6:
		{
			addrSize = sizeof(sockaddr_in6);
			memcpy(&addr6.sin6_addr.u, &addr.IPv6, sizeof(addr.IPv6));
			addr6.sin6_port = htons(port);
			socketAddr = (sockaddr*)&addr6;
		}
		break;
	}

	socketAddr->sa_family = (int)mFamily;
	WSA_CALL(::connect(mSocket, (const sockaddr*)socketAddr, addrSize));
}

void Socket::disconnect() {
	shutdown(mSocket, SD_BOTH);
	closesocket(mSocket);
}

void Socket::bind(uint16 port, const std::string& hostOrIp) {
	auto addr = resolveHost(hostOrIp);
	union
	{
		sockaddr_in addr4;
		sockaddr_in6 addr6;
	};
	sockaddr* socketAddr = NULL;
	uint32 addrSize = 0;
	switch(mFamily) {
	case IPv4: 
		{
			addrSize = sizeof(sockaddr_in);
			addr4.sin_addr.s_addr = addr.IPv4;
			addr4.sin_port = htons(port);
			socketAddr = (sockaddr*)&addr4;
		}
		break;
	case IPv6:
		{
			addrSize = sizeof(sockaddr_in6);
			memcpy(&addr6.sin6_addr.u, &addr.IPv6, sizeof(addr.IPv6));
			addr6.sin6_port = htons(port);
			socketAddr = (sockaddr*)&addr6;
		}
		break;
	}

	socketAddr->sa_family = (int)mFamily;
	WSA_CALL(::bind(mSocket, (const sockaddr*)socketAddr, addrSize));
}

void Socket::listen(uint16 backlog) {
	WSA_CALL(::listen(mSocket, backlog));
}

Socket* Socket::acceptSocket() {
	sockaddr_in sin;
	int len = sizeof(sin);
	SOCKET s = ::accept(mSocket, (sockaddr*)&sin, &len);
	if(s == INVALID_SOCKET)
		throw SocketException(WSAGetLastError(), "::accept(mSocket, NULL, NULL)");

	auto ret = new Socket(s, mType, mFamily, mProtocol);
	HostEntry pt;
	pt.IPv4 = sin.sin_addr.s_addr;
	pt.isIPv6 = false;
	pt.port = (sin.sin_port);
	ret->mRemotePoint = pt;
	return ret;
}

void Socket::send(const void* data, uint32 numBytes) {
	unsigned char* ptr = reinterpret_cast<unsigned char*>(const_cast<void*>(data));
	uint32 numSent = 0;
	while(numSent < numBytes) {
		int32 chunk = ::send(mSocket, (const char*)(ptr + numSent), numBytes - numSent, 0);
		if(chunk <= 0)
			throw SocketException(WSAGetLastError());

		numSent += chunk;
	}
}

void Socket::receive(void* buffer, uint32 numBytes) {
	unsigned char* ptr = reinterpret_cast<unsigned char*>(buffer);
	uint32 numRecv = 0;
	while(numRecv < numBytes) {
		int32 chunk = ::recv(mSocket, (char*)(ptr + numRecv), numBytes - numRecv, 0);
		if(chunk <= 0)
			throw SocketException(WSAGetLastError());

		numRecv += chunk;
	}
}

void Socket::setBlocking(bool enabled) {
	u_long val = enabled ? 1 : 0;
	ioctlsocket(mSocket, FIONBIO, &val);
}

uint32 Socket::getBytesAvailable() {
	u_long bytes = 0;
	ioctlsocket(mSocket, FIONREAD, &bytes);
	return bytes;
}

void Socket::setNagle(bool enabled) {
	BOOL val = enabled ? 0 : 1;
	setsockopt(mSocket, IPPROTO_TCP, TCP_NODELAY, (const char*)&val, sizeof(BOOL));
}